import torch
import torch.nn as nn
import torch.nn.functional as F

# 全连接层
# 首先我们要新建一个类，这个类要继承nn.Module


class LinearNet(nn.Module):
    """
    LinearNet
    """

    def __init__(self):
        super(LinearNet, self).__init__()
        self.fc2 = nn.Linear(1*28*28, 10)  # 最初传入的图片的像素点是1*28*28的，最后我们要收敛成10个结果
        # 如果先收敛成100个，然后在写一个全连接层
        # self.fc2 = nn.Linear(1 * 28 * 28, 100)
        # self.fc2 = nn.Linear(100, 10)
        # 激活函数，激励函数，通过数学手段将线性计算过程进行优化，使其加速。最常用的线性激活函数Relu
        self.relu = nn.ReLU()

    def forward(self, image):  # 继承前向传播的方法
        image_viwed = image.view(-1, 1*28*28)  # 此处需要拍平
        out = self.fc2(image_viwed)
        fcl_out = self.relu(out)  # 激活函数对应一下
        return out


class ConvNet(nn.Module):
    """ 
    ConvNet
    """

    def __init__(self):
        super().__init__()
        # batch*1*28*28（每次会送入batch个样本，输入通道数1（黑白图像），图像分辨率是28x28）
        # 下面的卷积层Conv2d的第一个参数指输入通道数，第二个参数指输出通道数，第三个参数指卷积核的大小
        self.conv1 = nn.Conv2d(1, 10, 5)  # 输入通道数1，输出通道数10，核的大小5
        self.conv2 = nn.Conv2d(10, 20, 3)  # 输入通道数10，输出通道数20，核的大小3
        # 下面的全连接层Linear的第一个参数指输入通道数，第二个参数指输出通道数
        self.fc1 = nn.Linear(20*10*10, 500)  # 输入通道数是2000，输出通道数是500
        self.fc2 = nn.Linear(500, 10)  # 输入通道数是500，输出通道数是10，即10分类

    def forward(self, x):
        # 在本例中in_size=512，也就是 BATCH_SIZE 的值。输入的x可以看成是512*1*28*28的张量。
        in_size = x.size(0)

        # batch*1*28*28 -> batch*10*24*24（28x28的图像经过一次核为5x5的卷积，输出变为24x24）
        out = self.conv1(x)
        out = F.relu(out)  # batch*10*24*24（激活函数ReLU不改变形状））

        # batch*10*24*24 -> batch*10*12*12（2*2的池化层会减半）
        out = F.max_pool2d(out, 2, 2)
        out = self.conv2(out)  # batch*10*12*12 -> batch*20*10*10（再卷积一次，核的大小是3）
        out = F.relu(out)  # batch*20*10*10

        # batch*20*10*10 -> batch*2000（out的第二维是-1，说明是自动推算，本例中第二维是20*10*10）
        out = out.view(in_size, -1)
        out = self.fc1(out)  # batch*2000 -> batch*500
        out = F.relu(out)  # batch*500
        out = self.fc2(out)  # batch*500 -> batch*10
        out = F.log_softmax(out, dim=1)  # 计算log(softmax(x))
        return out


class CNN(nn.Module):
    """
    CNN
    """

    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(
            # in_channels, out_channels, kernel_size, stride, padding
            nn.Conv2d(1, 16, 5, 1, 2),  # padding=2, 不改变形状 1*28*28 => 16*28*28
            nn.ReLU(),
            nn.MaxPool2d(2),  # 16*28*28 => 16*14*14
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),  # 16*14*14 => 32*14*14
            nn.ReLU(),
            nn.MaxPool2d(2),  # 32*14*14 => 32*7*7
        )
        # y = xA^T + b
        # 输入的神经元个数 = 7*7*32 表示 32 channels, 每 channel 为 7*7
        self.fc = nn.Linear(7*7*32, 10)

    def forward(self, x):
        x = self.conv1(x)   # 16*14*14
        x = self.conv2(x)   # 32*7*7
        x = torch.Tensor(x).view(x.size(0), -1)  # reshape => ?
        out = self.fc(x)
        return out, x
